#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <ctype.h>
#include <float.h> // For DBL_MAX


// (1) Read PGM file into matrix
void pgm_to_matrix(const char* file_path, const int n_rows, const int n_cols, double A[n_rows][n_cols]) {
    FILE* fp = fopen(file_path, "rb");
    if (!fp) {
        perror("Error opening file");
        exit(EXIT_FAILURE);
    }

    // Skip first 4 header lines
    char buffer[256];
    for (int i = 0; i < 4; i++) {
        if (!fgets(buffer, sizeof(buffer), fp)) {
            fclose(fp);
            fprintf(stderr, "Invalid PGM header\n");
            exit(EXIT_FAILURE);
        }
    }

    // Read pixel data and normalize
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            int pixel = fgetc(fp);
            if (pixel == EOF) {
                fclose(fp);
                fprintf(stderr, "Unexpected end of file\n");
                exit(EXIT_FAILURE);
            }
            A[i][j] = (double)pixel / 255.0;  // Normalize to [0,1]
        }
    }
    fclose(fp);
}

// (2) Flatten matrix to 1D vector
void matrix_to_vector(const int n_rows, const int n_cols, 
                      const double matrix[n_rows][n_cols], 
                      double vector[1][n_rows * n_cols]) {
    int idx = 0;
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            vector[0][idx++] = matrix[i][j];
        }
    }
}

// (3) Reshape 1D vector to 2D matrix
void vector_to_matrix(const int n_rows, const int n_cols,
                      const double vector[1][n_rows * n_cols],
                      double matrix[n_rows][n_cols]) {
    int idx = 0;
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            matrix[i][j] = vector[0][idx++];
        }
    }
}

// (4) Read numeric data from text file into 1D vector
void read_txt_data(const char* filename, const int n, double array[1][n]) {
    FILE* fp = fopen(filename, "r");
    if (!fp) {
        perror("Error opening weight file");
        exit(EXIT_FAILURE);
    }

    for (int i = 0; i < n; i++) {
        if (fscanf(fp, "%lf", &array[0][i]) != 1) {
            fclose(fp);
            fprintf(stderr, "Error reading weight data at position %d\n", i);
            exit(EXIT_FAILURE);
        }
    }
    fclose(fp);
}

// (5) Element-wise matrix addition
void add_matrices(const int n_rows, const int n_cols, 
                  const double mat1[n_rows][n_cols],
                  const double mat2[n_rows][n_cols], 
                  double result[n_rows][n_cols]) {
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            result[i][j] = mat1[i][j] + mat2[i][j];
        }
    }
}

// (6) Matrix multiplication
void matrix_mul(const int n_rows, const int n_cols, const int p,
                const double matrix1[n_rows][p],
                const double matrix2[p][n_cols],
                double matrix_product[n_rows][n_cols]) {
    // Initialize result matrix to zero
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            matrix_product[i][j] = 0.0;
        }
    }
    
    // Perform matrix multiplication
    for (int i = 0; i < n_rows; i++) {
        for (int k = 0; k < p; k++) {
            for (int j = 0; j < n_cols; j++) {
                matrix_product[i][j] += matrix1[i][k] * matrix2[k][j];
            }
        }
    }
}

// (7) ReLU activation function
void active_function(const int SIZE, const double F[1][SIZE], double G[1][SIZE]) {
    for (int i = 0; i < SIZE; i++) {
        G[0][i] = (F[0][i] > 0) ? F[0][i] : 0.0;  // max(F, 0)
    }
}

// (8) Softmax function for probability distribution
void softmax(const int n_rows, const int n_cols, 
             const double L[n_rows][n_cols], 
             double S[n_rows][n_cols]) {
    // Find maximum value in matrix for numerical stability
    double max_val = -DBL_MAX;
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            if (L[i][j] > max_val) max_val = L[i][j];
        }
    }
    
    // Calculate exponential sum
    double sum_exp = 0.0;
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            sum_exp += exp(L[i][j] - max_val);
        }
    }
    
    // Compute SoftMax probabilities
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            S[i][j] = exp(L[i][j] - max_val) / sum_exp;
        }
    }
}

// ====================== Phase II: Digit Recognition Pipeline ======================

// Recognize digit in single PGM file
int recognize_digit(const char* filename) {
    // Step 1: Read and flatten image
    double A[28][28];      // 28x28 image matrix
    double B[1][784];      // Flattened 1x784 vector
    
    pgm_to_matrix(filename, 28, 28, A);
    matrix_to_vector(28, 28, A, B);
    
    // Step 2: First layer weights (B × W1)
    double W1_vector[1][100352]; // Raw data from W1.txt
    read_txt_data("input_files/W1.txt", 100352, W1_vector);
    
    double W1_matrix[784][128];  // Reshaped 784x128 weight matrix
    vector_to_matrix(784, 128, W1_vector, W1_matrix);
    
    double D[1][128];            // First layer output
    matrix_mul(1, 128, 784, B, W1_matrix, D);
    
    // Step 3: Add first layer bias (D + B1)
    double E[1][128];            // Bias vector from B1.txt
    read_txt_data("input_files/B1.txt", 128, E);
    
    double F[1][128];            // Result after bias
    add_matrices(1, 128, D, E, F);
    
    // Step 4: Apply ReLU activation
    double G[1][128];            // Activated output
    active_function(128, F, G);
    
    // Step 5: Second layer weights (G × W2)
    double W2_vector[1][1280];   // Raw data from W2.txt
    read_txt_data("input_files/W2.txt", 1280, W2_vector);
    
    double W2_matrix[128][10];   // Reshaped 128x10 weight matrix
    vector_to_matrix(128, 10, W2_vector, W2_matrix);
    
    double H[1][10];             // Second layer output
    matrix_mul(1, 10, 128, G, W2_matrix, H);
    
    // Step 6: Add second layer bias (H + B2)
    double J[1][10];             // Bias vector from B2.txt
    read_txt_data("input_files/B2.txt", 10, J);
    
    double L[1][10];             // Final output layer
    add_matrices(1, 10, H, J, L);
    
    // Step 7: Apply SoftMax for probabilities
    double S[1][10];             // Probability distribution
    softmax(1, 10, L, S);
    
    // Determine digit with highest probability
    int digit = 0;
    double max_prob = S[0][0];
    for (int i = 1; i < 10; i++) {
        if (S[0][i] > max_prob) {
            max_prob = S[0][i];
            digit = i;
        }
    }
    
    return digit;
}


// Check if file has .pgm extension
int is_pgm_file(const char* filename) {
    const char* dot = strrchr(filename, '.');
    return dot && !strcasecmp(dot, ".pgm");
}

int main() {
    char input[1024];
    char filenames[10][256]; // Supports up to 10 files
    int file_count;
    
    while (1) {
        printf("please input filenames: ");
        if (!fgets(input, sizeof(input), stdin)) break;
        
        // Remove trailing newline
        input[strcspn(input, "\n")] = 0;
        
        // Parse filenames separated by spaces or commas
        file_count = 0;
        char* token = strtok(input, " ,");
        while (token && file_count < 10) {
            strncpy(filenames[file_count], token, sizeof(filenames[0]) - 1);
            filenames[file_count][sizeof(filenames[0]) - 1] = '\0';
            file_count++;
            token = strtok(NULL, " ,");
        }
        
        // Process each file
        int valid_files = 0;
        for (int i = 0; i < file_count; i++) {
            // Validate file extension
            if (!is_pgm_file(filenames[i])) {
                printf("invalid file: %s\n", filenames[i]);
                continue;
            }
            
            // Check file existence
            FILE* fp = fopen(filenames[i], "rb");
            if (!fp) {
                printf("invalid file: %s\n", filenames[i]);
                continue;
            }
            fclose(fp);
            
            // Recognize digit and print result
            int digit = recognize_digit(filenames[i]);
            printf("%s:%d\n", filenames[i], digit);
            valid_files++;
        }
        
        // Handle case with no valid files
        if (valid_files == 0 && file_count > 0) {
            printf("invalid file\n");
        }
        
        // Continue prompt
        while (1) {
            printf("do you want to continue? please input[Y or N]:");
            if (!fgets(input, sizeof(input), stdin)) break;
            
            // Normalize input to uppercase
            input[strcspn(input, "\n")] = 0;
            for (char* p = input; *p; p++) *p = toupper(*p);
            
            if (strcmp(input, "Y") == 0) {
                break; // Continue processing
            } else if (strcmp(input, "N") == 0) {
                printf("Bye\n");
                return 0;
            } else {
                printf("Invalid input. Please enter Y or N.\n");
            }
        }
    }
    
    return 0;
}